package arithmetic3;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Stack;

/**
 * 并查集结构的实现
 */
public class UnionFind<T> {
    //存储一个元素和它对应的包装节点的集合
    private Map<T, Node<T>> nodeMap = new HashMap<>();
    //存储一个节点和它父节点的集合
    private Map<Node<T>, Node<T>> fatherMap = new HashMap<>();
    //存储一个节点和它所代表的集合的大小的集合
    private Map<Node<T>, Integer> sizeMap = new HashMap<>();

    /**
     * 初始化并查集
     */
    public UnionFind(List<T> list) {
        if (list != null) {
            for (T t : list) {
                Node<T> node = new Node<T>(t);
                nodeMap.put(t, node);
                //初始时每个节点的父节点为自己
                fatherMap.put(node, node);
                sizeMap.put(node, 1);
            }
        }
    }

    /**
     * 得到一个节点的代表节点
     *
     * @param node
     * @return
     */
    public Node<T> findFather(Node<T> node) {
        //准备一个栈，在寻找父节点的同时优化并查集的结构
        Stack<Node<T>> stack = new Stack<>();
        while (node != fatherMap.get(node)) {
            stack.push(fatherMap.get(node));
            node = fatherMap.get(node);
        }
        //将栈中的所有节点的代表节点变为node
        while (!stack.isEmpty()) {
            fatherMap.put(stack.pop(), node);
        }
        return node;
    }

    /**
     * 查询两个元素在不在一个集合中
     *
     * @return
     */
    public boolean isSame(T t1, T t2) {
        //两个节点都存时
        if (nodeMap.containsKey(t1) && nodeMap.containsKey(t2)) {
            return fatherMap.get(nodeMap.get(t1)) == fatherMap.get(nodeMap.get(t2));
        }
        return false;
    }

    /**
     * 将两个元素所在的集合合并为1个集合
     *
     * @param t1
     * @param t2
     */
    public void union(T t1, T t2) {
        //当两个节点都存在，并且不再一个集合中时才合并
        if (nodeMap.containsKey(t1) && nodeMap.containsKey(t2) && !isSame(t1, t2)) {
            Node<T> node1 = nodeMap.get(t1);
            Node<T> node2 = nodeMap.get(t2);
            //得到两个节点代表节点
            Node<T> father1 = findFather(node1);
            Node<T> father2 = findFather(node2);
            //将较小的集合合并到较大的集合里
            Integer size1 = sizeMap.get(father1);
            Integer size2 = sizeMap.get(father2);
            Node<T> bigFather = size1 > size2 ? father1 : father2;
            Node<T> smallFather = size1 < size2 ? father1 : father2;
            fatherMap.put(smallFather, bigFather);
            sizeMap.remove(smallFather);
            sizeMap.put(bigFather, size1 + size2);
        }
    }

    /**
     * 包装类型
     */
    private static class Node<T> {
        private T value;

        public Node(T value) {
            this.value = value;
        }
    }
}
